from django.db import models
from django.db import transaction


class RequestQuerySetError(BaseException):
    pass


class RequestQuerySet(models.QuerySet):
    def __init__(self, model=None, query=None, using=None, hints=None):
        super(RequestQuerySet, self).__init__(model, query, using, hints)
        self.request = None

    def of(self, request):
        self.request = request
        return self
    
    def _use_super_func(self):
        if self.request is None:
            return True
        if not hasattr(self.request, 'user'):
            return True
        if not hasattr(self.request.user, 'global_user'):
            return True
        return self.request.user.global_user


    def count(self):
        return super(RequestQuerySet, self).count()

    def get(self, *args, **kwargs):
        return super(RequestQuerySet, self).get(*args, **kwargs)

    def create(self, **kwargs):
        if self._use_super_func():
            return super(RequestQuerySet, self).create(**kwargs)
        return super(RequestQuerySet, self).create(firm_id=self.request.user.firm_id, **kwargs)

    def bulk_create(self, objs, batch_size=None, ignore_conflicts=False):
        if self._use_super_func():
            return super(RequestQuerySet, self).bulk_create(objs, batch_size, ignore_conflicts)

        for o in objs:
            o.firm_id = self.request.user.firm_id
        return super(RequestQuerySet, self).bulk_create(objs, batch_size, ignore_conflicts)

    def get_or_create(self, defaults=None, **kwargs):
        if self._use_super_func():
            return super(RequestQuerySet, self).get_or_create(defaults, **kwargs)

        self._for_write = True
        try:
            return self.get(**kwargs), False
        except self.model.DoesNotExist:
            params = self._extract_model_params(defaults, firm_id=self.request.user.firm_id, **kwargs)
            return self._create_object_from_params(kwargs, params)

    def update_or_create(self, defaults=None, **kwargs):
        if self._use_super_func():
            return super(RequestQuerySet, self).update_or_create(defaults, **kwargs)
        defaults = defaults or {}
        self._for_write = True
        with transaction.atomic(using=self.db):
            try:
                obj = self.select_for_update().get(**kwargs)
            except self.model.DoesNotExist:
                params = self._extract_model_params(defaults, **kwargs)
                # Lock the row so that a concurrent update is blocked until
                # after update_or_create() has performed its save.
                obj, created = self._create_object_from_params(kwargs, params, lock=True)
                if created:
                    return obj, created
            for k, v in defaults.items():
                setattr(obj, k, v() if callable(v) else v)
            obj.firm_id = self.request.user.firm_id
            obj.save(using=self.db)
        return obj, False

    def filter(self, *args, **kwargs):
        if self._use_super_func():
            return super(RequestQuerySet, self).filter(*args, **kwargs)
        return super(RequestQuerySet, self).filter(firm_id=self.request.user.firm_id, *args, **kwargs)

    def all(self):
        if self._use_super_func():
            return super(RequestQuerySet, self).all()
        return super(RequestQuerySet, self).filter(firm_id=self.request.user.firm_id)

